import time
import os.path as osp
from tensorboardX import SummaryWriter
class StepLog(object):
    def __init__(self,ep,num_models,step,step_st,log_steps,
                    g_loss_weight,s_loss_weight,id_loss_weight,
                    pm_loss_weight,gdm_loss_weight,sdm_loss_weight,
                    g_prec_meter,g_m_meter,s_prec_meter,s_m_meter,
                    g_dist_ap_meter,g_dist_an_meter,s_dist_ap_meter,s_dist_an_meter,
                    loss_meter,g_loss_meter,s_loss_meter,id_loss_meter,
                    pm_loss_meter,gdm_loss_meter,sdm_loss_meter):
        ############
        # Step Log #
        ############

        if step % log_steps == 0:
            time_log = '\tStep {}/Ep {}, {:.2f}s'.format(
                step, ep + 1, time.time() - step_st, )

            if g_loss_weight > 0:
                g_log = (', gp {:.2%}, gm {:.2%}, '
                         'gd_ap {:.4f}, gd_an {:.4f}, '
                         'gL {:.4f}'.format(
                    g_prec_meter.val, g_m_meter.val,
                    g_dist_ap_meter.val, g_dist_an_meter.val,
                    g_loss_meter.val, ))
            else:
                g_log = ''

            if s_loss_weight > 0:
                l_log = (', sp {:.2%}, sm {:.2%}, '
                         'sd_ap {:.4f}, sd_an {:.4f}, '
                         'sL {:.4f}'.format(
                    s_prec_meter.val, s_m_meter.val,
                    s_dist_ap_meter.val, s_dist_an_meter.val,
                    s_loss_meter.val, ))
            else:
                l_log = ''

            if id_loss_weight > 0:
                id_log = (', idL {:.4f}'.format(id_loss_meter.val))
            else:
                id_log = ''

            if (num_models > 1) and (pm_loss_weight > 0):
                pm_log = (', pmL {:.4f}'.format(pm_loss_meter.val))
            else:
                pm_log = ''

            if (num_models > 1) and (gdm_loss_weight > 0):
                gdm_log = (', gdmL {:.4f}'.format(gdm_loss_meter.val))
            else:
                gdm_log = ''

            if (num_models > 1) \
                    and (sdm_loss_weight > 0):
                ldm_log = (', sdmL {:.4f}'.format(sdm_loss_meter.val))
            else:
                ldm_log = ''

            total_loss_log = ', loss {:.4f}'.format(loss_meter.val)

            log = time_log + \
                  g_log + l_log + id_log + \
                  pm_log + gdm_log + ldm_log + \
                  total_loss_log
            print(log)

class EpochLog(object):
    def __init__(self,ep,ep_st,num_models,g_loss_weight,s_loss_weight,id_loss_weight,
                 pm_loss_weight,gdm_loss_weight,sdm_loss_weight,
                 g_prec_meter,g_m_meter,s_prec_meter,s_m_meter,
                 g_dist_ap_meter,g_dist_an_meter,
                 s_dist_ap_meter,s_dist_an_meter,
                 loss_meter,g_loss_meter,s_loss_meter,id_loss_meter,
                 pm_loss_meter,gdm_loss_meter,sdm_loss_meter):
        #############
        # Epoch Log #
        #############

        time_log = 'Ep {}, {:.2f}s'.format(ep + 1, time.time() - ep_st, )

        if g_loss_weight > 0:
            g_log = (', gp {:.2%}, gm {:.2%}, '
                     'gd_ap {:.4f}, gd_an {:.4f}, '
                     'gL {:.4f}'.format(
                g_prec_meter.avg, g_m_meter.avg,
                g_dist_ap_meter.avg, g_dist_an_meter.avg,
                g_loss_meter.avg, ))
        else:
            g_log = ''

        if s_loss_weight > 0:
            l_log = (', sp {:.2%}, sm {:.2%}, '
                     'sd_ap {:.4f}, sd_an {:.4f}, '
                     'sL {:.4f}'.format(
                s_prec_meter.avg, s_m_meter.avg,
                s_dist_ap_meter.avg, s_dist_an_meter.avg,
                s_loss_meter.avg, ))
        else:
            l_log = ''

        if id_loss_weight > 0:
            id_log = (', idL {:.4f}'.format(id_loss_meter.avg))
        else:
            id_log = ''

        if (num_models > 1) and (pm_loss_weight > 0):
            pm_log = (', pmL {:.4f}'.format(pm_loss_meter.avg))
        else:
            pm_log = ''

        if (num_models > 1) and (gdm_loss_weight > 0):
            gdm_log = (', gdmL {:.4f}'.format(gdm_loss_meter.avg))
        else:
            gdm_log = ''

        if (num_models > 1) \
                and (sdm_loss_weight > 0):
            ldm_log = (', sdmL {:.4f}'.format(sdm_loss_meter.avg))
        else:
            ldm_log = ''

        total_loss_log = ', loss {:.4f}'.format(loss_meter.avg)

        log = time_log + \
              g_log + l_log + id_log + \
              pm_log + gdm_log + ldm_log + \
              total_loss_log
        print(log)

class Log2TensorBoard(object):
    def __init__(self,ep,log_to_file,exp_dir,writer,loss_meter,g_loss_meter,s_loss_meter,id_loss_meter,
                 g_m_meter, s_m_meter, g_dist_ap_meter, g_dist_an_meter, s_dist_ap_meter, s_dist_an_meter,
                 pm_loss_meter,gdm_loss_meter,sdm_loss_meter,g_prec_meter,s_prec_meter,):
        if log_to_file:
            if writer is None:
                writer = SummaryWriter(log_dir=osp.join(exp_dir, 'tensorboard'))
            writer.add_scalars(
                'loss',
                dict(global_loss=g_loss_meter.avg,
                     local_loss=s_loss_meter.avg,
                     id_loss=id_loss_meter.avg,
                     pm_loss=pm_loss_meter.avg,
                     gdm_loss=gdm_loss_meter.avg,
                     ldm_loss=sdm_loss_meter.avg,
                     loss=loss_meter.avg, ),
                ep)
            writer.add_scalars(
                'tri_precision',
                dict(global_precision=g_prec_meter.avg,
                     local_precision=s_prec_meter.avg, ),
                ep)
            writer.add_scalars(
                'satisfy_margin',
                dict(global_satisfy_margin=g_m_meter.avg,
                     local_satisfy_margin=s_m_meter.avg, ),
                ep)
            writer.add_scalars(
                'global_dist',
                dict(global_dist_ap=g_dist_ap_meter.avg,
                     global_dist_an=g_dist_an_meter.avg, ),
                ep)
            writer.add_scalars(
                'local_dist',
                dict(local_dist_ap=s_dist_ap_meter.avg,
                     local_dist_an=s_dist_an_meter.avg, ),
                ep)